#ifndef _YOLOLAYER_DYNAMIC_PLUGIN_H
#define _YOLOLAYER_DYNAMIC_PLUGIN_H

#include <vector>
#include <string>
#include "NvInfer.h"
#include "yololayer.h"

namespace nvinfer1
{
    class YoloLayerDynamicPlugin: public IPluginV2DynamicExt
    {
    public:
        YoloLayerDynamicPlugin(int classCount, int maxOut, const std::vector<Yolo::YoloKernel>& vYoloKernel);
        YoloLayerDynamicPlugin(const void* data, size_t length);
        ~YoloLayerDynamicPlugin();

        //
        const char* getPluginType() const override;
        const char* getPluginVersion() const override;
        int getNbOutputs() const override;
        int initialize() override;
        void terminate() override {}
        size_t getSerializationSize() const override;
        void serialize(void* buffer) const override;
        void destroy() override;
        void setPluginNamespace(const char* libNamespace) override;
        const char* getPluginNamespace() const override;

        // IPluginV2Ext methods
        DataType getOutputDataType(int index, const nvinfer1::DataType* inputType, int nbInputs) const override;

        // IPluginV2DynamicExt methods
        IPluginV2DynamicExt* clone() const override;
        DimsExprs getOutputDimensions(
           int outputIndex, const DimsExprs* inputs, int nbInputs, IExprBuilder& exprBuilder) override;
        bool supportsFormatCombination(
           int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) override;
        void configurePlugin(const DynamicPluginTensorDesc* in, int nbInputs, const DynamicPluginTensorDesc* out,
           int nbOutputs) override;
        size_t getWorkspaceSize(const PluginTensorDesc* inputs, int nbInputs, const PluginTensorDesc* outputs,
            int nbOutputs) const override;
        int enqueue(const PluginTensorDesc* inputDesc, const PluginTensorDesc* outputDesc, const void* const* inputs,
           void* const* outputs, void* workspace, cudaStream_t stream) override;

    private:
//        void forwardGpu(const float* const* inputs, float *output, cudaStream_t stream, int batchSize);
        void forwardGpu(nvinfer1::DataType data_type, const void* const* inputs, float *output, cudaStream_t stream, int batchSize);
        int mThreadCount = 256;
        int mKernelCount;
        int mClassCount;
        int mYoloV5NetWidth;
        int mYoloV5NetHeight;
        int mMaxOutObject;
        std::vector<Yolo::YoloKernel> mYoloKernel;
        void** mAnchor;
        const char* mPluginNamespace;
    };

    class YoloLayerDynamicPluginCreator : public IPluginCreator
    {
    public:
        YoloLayerDynamicPluginCreator();
        ~YoloLayerDynamicPluginCreator() override;

        const char* getPluginName() const override;

        const char* getPluginVersion() const override;

        const PluginFieldCollection* getFieldNames() override;

        IPluginV2DynamicExt* createPlugin(const char* name, const PluginFieldCollection* fc) override;

        IPluginV2DynamicExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override;

        void setPluginNamespace(const char* libNamespace) override
        {
            mNamespace = libNamespace;
        }

        const char* getPluginNamespace() const override
        {
            return mNamespace.c_str();
        }

    private:
        std::string mNamespace;
        static PluginFieldCollection mFC;
        static std::vector<PluginField> mPluginAttributes;
    };
    REGISTER_TENSORRT_PLUGIN(YoloLayerDynamicPluginCreator);
};

#endif 
